/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.solr.handler.component;

import java.util.Arrays;
import java.util.Iterator;
import java.util.Map;
import org.apache.solr.BaseDistributedSearchTestCase;
import org.apache.solr.client.solrj.response.QueryResponse;
import org.apache.solr.common.SolrDocument;
import org.apache.solr.common.SolrDocumentList;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.params.ModifiableSolrParams;
import org.junit.BeforeClass;

/**
 * Test for distributed ExpandComponent
 *
 * @see org.apache.solr.handler.component.ExpandComponent
 */
public class DistributedExpandComponentTest extends BaseDistributedSearchTestCase {

  public DistributedExpandComponentTest() {
    stress = 0;
  }

  @BeforeClass
  public static void setUpBeforeClass() throws Exception {
    initCore("solrconfig-collapseqparser.xml", "schema11.xml");
  }

  @ShardsFixed(num = 3)
  public void test() throws Exception {
    _test("group_s", "g1", "g2", "g3", "g4");
    _test("group_s_dv", "g1", "g2", "g3", "g4");
    _test(
        "group_i", "1", "0", "3", "-1"); // NOTE: using 0 to explicitly confim we don't assume null
    _test(
        "group_ti_dv",
        "1",
        "-2",
        "0",
        "4"); // NOTE: using 0 to explicitly confim we don't assume null
  }

  private void _test(
      final String group, final String aaa, final String bbb, final String ccc, final String ddd)
      throws Exception {

    del("*:*");

    index_specific(
        0, "id", "1", "term_s", "YYYY", group, aaa, "test_i", "5", "test_l", "10", "test_f",
        "2000");
    index_specific(
        0, "id", "2", "term_s", "YYYY", group, aaa, "test_i", "50", "test_l", "100", "test_f",
        "200");
    index_specific(
        1, "id", "5", "term_s", "YYYY", group, bbb, "test_i", "4", "test_l", "10", "test_f",
        "2000");
    index_specific(
        1, "id", "6", "term_s", "YYYY", group, bbb, "test_i", "10", "test_l", "100", "test_f",
        "200");
    index_specific(
        0, "id", "7", "term_s", "YYYY", group, aaa, "test_i", "1", "test_l", "100000", "test_f",
        "2000");
    index_specific(
        1, "id", "8", "term_s", "YYYY", group, bbb, "test_i", "2", "test_l", "100000", "test_f",
        "200");
    index_specific(
        2, "id", "9", "term_s", "YYYY", group, ccc, "test_i", "1000", "test_l", "1005", "test_f",
        "3000");
    index_specific(
        2, "id", "10", "term_s", "YYYY", group, ccc, "test_i", "1500", "test_l", "1001", "test_f",
        "3200");

    // NOTE: nullPolicy=collapse will only be viable because all null docs are in colocated in
    // shard #2
    index_specific(2, "id", "88", "test_i", "1001", "test_l", "1001", "test_f", "3200");
    index_specific(2, "id", "99", "test_i", "11", "test_l", "100", "test_f", "200");

    index_specific(
        2, "id", "11", "term_s", "YYYY", group, ccc, "test_i", "1300", "test_l", "1002", "test_f",
        "3300");
    index_specific(
        1, "id", "12", "term_s", "YYYY", group, ddd, "test_i", "15", "test_l", "10", "test_f",
        "2000");
    index_specific(
        1, "id", "13", "term_s", "YYYY", group, ddd, "test_i", "16", "test_l", "9", "test_f",
        "2000");
    index_specific(
        1, "id", "14", "term_s", "YYYY", group, ddd, "test_i", "1", "test_l", "20", "test_f",
        "2000");

    commit();

    handle.put("explain", SKIPVAL);
    handle.put("timestamp", SKIPVAL);
    handle.put("score", SKIPVAL);
    handle.put("wt", SKIP);
    handle.put("distrib", SKIP);
    handle.put("shards.qt", SKIP);
    handle.put("shards", SKIP);
    handle.put("q", SKIP);
    handle.put("maxScore", SKIPVAL);
    handle.put("_version_", SKIP);
    handle.put("expanded", UNORDERED);

    query(
        "q",
        "*:*",
        "fq",
        "{!collapse field=" + group + "}",
        "defType",
        "edismax",
        "bf",
        "field(test_i)",
        "expand",
        "true",
        "fl",
        "*,score");
    query(
        "q",
        "*:*",
        "fq",
        "{!collapse field=" + group + "}",
        "defType",
        "edismax",
        "bf",
        "field(test_i)",
        "expand",
        "true",
        "expand.sort",
        "test_l desc",
        "fl",
        "*,score");
    query(
        "q",
        "*:*",
        "fq",
        "{!collapse field=" + group + "}",
        "defType",
        "edismax",
        "bf",
        "field(test_i)",
        "expand",
        "true",
        "expand.sort",
        "test_l desc",
        "expand.rows",
        "1",
        "fl",
        "*,score");
    // Test no expand results
    query(
        "q",
        "test_i:5",
        "fq",
        "{!collapse field=" + group + "}",
        "defType",
        "edismax",
        "bf",
        "field(test_i)",
        "expand",
        "true",
        "expand.sort",
        "test_l desc",
        "expand.rows",
        "1",
        "fl",
        "*,score");
    // Test zero results
    query(
        "q",
        "test_i:5434343",
        "fq",
        "{!collapse field=" + group + "}",
        "defType",
        "edismax",
        "bf",
        "field(test_i)",
        "expand",
        "true",
        "expand.sort",
        "test_l desc",
        "expand.rows",
        "1",
        "fl",
        "*,score");
    // Test page 2
    query(
        "q",
        "*:*",
        "start",
        "1",
        "rows",
        "1",
        "fq",
        "{!collapse field=" + group + "}",
        "defType",
        "edismax",
        "bf",
        "field(test_i)",
        "expand",
        "true",
        "fl",
        "*,score");

    // multiple collapse and equal cost
    ModifiableSolrParams baseParams =
        params(
            "q",
            "term_s:YYYY",
            "defType",
            "edismax",
            "expand",
            "true",
            "fl",
            "*,score",
            "bf",
            "field(test_i)",
            "expand.sort",
            "id asc");
    baseParams.set("fq", "{!collapse field=" + group + "}", "{!collapse field=test_i}");
    query(baseParams);

    // multiple collapse and unequal cost case1
    baseParams.set(
        "fq", "{!collapse cost=1000 field=" + group + "}", "{!collapse cost=2000 field=test_i}");
    query(baseParams);

    // multiple collapse and unequal cost case2
    baseParams.set(
        "fq", "{!collapse cost=1000 field=" + group + "}", "{!collapse cost=200 field=test_i}");
    query(baseParams);

    ignoreException("missing expand field");
    SolrException e = expectThrows(SolrException.class, () -> query("q", "*:*", "expand", "true"));
    assertEquals(SolrException.ErrorCode.BAD_REQUEST.code, e.code());
    assertTrue(e.getMessage().contains("missing expand field"));
    resetExceptionIgnores();

    // Since none of these queries will match any doc w/null in the group field, it shouldn't matter
    // what nullPolicy is used...
    for (String np :
        Arrays.asList("", " nullPolicy=ignore", " nullPolicy=expand", " nullPolicy=collapse")) {

      // First basic test case.
      ModifiableSolrParams params = new ModifiableSolrParams();
      params.add("q", "term_s:YYYY");
      params.add("fq", "{!collapse field=" + group + np + "}");
      params.add("defType", "edismax");
      params.add("bf", "field(test_i)");
      params.add("expand", "true");

      setDistributedParams(params);
      QueryResponse rsp = query(params);
      assertCountAndOrder(
          4, rsp.getResults(), "10" /* c */, "2" /* a */, "13" /* d */, "6" /* b */);
      Map<String, SolrDocumentList> results = rsp.getExpandedResults();
      assertExpandGroups(results, aaa, bbb, ccc, ddd);
      assertExpandGroupCountAndOrder(aaa, 2, results, "1", "7");
      assertExpandGroupCountAndOrder(bbb, 2, results, "5", "8");
      assertExpandGroupCountAndOrder(ccc, 2, results, "11", "9");
      assertExpandGroupCountAndOrder(ddd, 2, results, "12", "14");

      // Test expand.sort

      params = new ModifiableSolrParams();
      params.add("q", "term_s:YYYY");
      params.add("fq", "{!collapse field=" + group + np + "}");
      params.add("defType", "edismax");
      params.add("bf", "field(test_i)");
      params.add("expand", "true");
      params.add("expand.sort", "test_l desc");
      setDistributedParams(params);
      rsp = query(params);
      assertCountAndOrder(
          4, rsp.getResults(), "10" /* c */, "2" /* a */, "13" /* d */, "6" /* b */);
      results = rsp.getExpandedResults();
      assertExpandGroups(results, aaa, bbb, ccc, ddd);
      assertExpandGroupCountAndOrder(aaa, 2, results, "7", "1");
      assertExpandGroupCountAndOrder(bbb, 2, results, "8", "5");
      assertExpandGroupCountAndOrder(ccc, 2, results, "9", "11");
      assertExpandGroupCountAndOrder(ddd, 2, results, "14", "12");

      // Test expand.rows

      params = new ModifiableSolrParams();
      params.add("q", "term_s:YYYY");
      params.add("fq", "{!collapse field=" + group + np + "}");
      params.add("defType", "edismax");
      params.add("bf", "field(test_i)");
      params.add("expand", "true");
      params.add("expand.sort", "test_l desc");
      params.add("expand.rows", "1");
      setDistributedParams(params);
      rsp = query(params);
      assertCountAndOrder(
          4, rsp.getResults(), "10" /* c */, "2" /* a */, "13" /* d */, "6" /* b */);
      results = rsp.getExpandedResults();
      assertExpandGroups(results, aaa, bbb, ccc, ddd);
      assertExpandGroupCountAndOrder(aaa, 1, results, "7");
      assertExpandGroupCountAndOrder(bbb, 1, results, "8");
      assertExpandGroupCountAndOrder(ccc, 1, results, "9");
      assertExpandGroupCountAndOrder(ddd, 1, results, "14");

      // Test expand.rows = 0 - no docs only expand count

      params = new ModifiableSolrParams();
      params.add("q", "term_s:YYYY");
      params.add("fq", "{!collapse field=" + group + np + "}");
      params.add("defType", "edismax");
      params.add("bf", "field(test_i)");
      params.add("expand", "true");
      params.add("expand.rows", "0");
      params.add("fl", "id");
      setDistributedParams(params);
      rsp = query(params);
      assertCountAndOrder(
          4, rsp.getResults(), "10" /* c */, "2" /* a */, "13" /* d */, "6" /* b */);
      results = rsp.getExpandedResults();
      assertExpandGroups(results, aaa, bbb, ccc, ddd);
      assertExpandGroupCountAndOrder(aaa, 0, results);
      assertExpandGroupCountAndOrder(bbb, 0, results);
      assertExpandGroupCountAndOrder(ccc, 0, results);
      assertExpandGroupCountAndOrder(ddd, 0, results);

      // Test expand.rows = 0 with expand.field

      params = new ModifiableSolrParams();
      params.add("q", "term_s:YYYY");
      params.add("fq", "test_l:10");
      params.add("defType", "edismax");
      params.add("bf", "field(test_i)");
      params.add("expand", "true");
      params.add("expand.fq", "test_f:2000");
      params.add("expand.field", group);
      params.add("expand.rows", "0");
      params.add("fl", "id,score");
      setDistributedParams(params);
      rsp = query(params);
      assertCountAndOrder(3, rsp.getResults(), "12" /* d */, "1" /* a */, "5" /* b */);
      results = rsp.getExpandedResults();
      assertExpandGroups(results, aaa, ddd);
      assertExpandGroupCountAndOrder(aaa, 0, results);
      assertExpandGroupCountAndOrder(ddd, 0, results);

      // Test key-only fl

      params = new ModifiableSolrParams();
      params.add("q", "term_s:YYYY");
      params.add("fq", "{!collapse field=" + group + np + "}");
      params.add("defType", "edismax");
      params.add("bf", "field(test_i)");
      params.add("expand", "true");
      params.add("fl", "id");

      setDistributedParams(params);
      rsp = query(params);
      assertCountAndOrder(
          4, rsp.getResults(), "10" /* c */, "2" /* a */, "13" /* d */, "6" /* b */);
      results = rsp.getExpandedResults();
      assertExpandGroups(results, aaa, bbb, ccc, ddd);
      assertExpandGroupCountAndOrder(aaa, 2, results, "1", "7");
      assertExpandGroupCountAndOrder(bbb, 2, results, "5", "8");
      assertExpandGroupCountAndOrder(ccc, 2, results, "11", "9");
      assertExpandGroupCountAndOrder(ddd, 2, results, "12", "14");

      // Test distrib.singlePass true

      params = new ModifiableSolrParams();
      params.add("q", "term_s:YYYY");
      params.add("fq", "{!collapse field=" + group + np + "}");
      params.add("defType", "edismax");
      params.add("bf", "field(test_i)");
      params.add("expand", "true");
      params.add("distrib.singlePass", "true");

      setDistributedParams(params);
      rsp = query(params);
      assertCountAndOrder(
          4, rsp.getResults(), "10" /* c */, "2" /* a */, "13" /* d */, "6" /* b */);
      results = rsp.getExpandedResults();
      assertExpandGroups(results, aaa, bbb, ccc, ddd);
      assertExpandGroupCountAndOrder(aaa, 2, results, "1", "7");
      assertExpandGroupCountAndOrder(bbb, 2, results, "5", "8");
      assertExpandGroupCountAndOrder(ccc, 2, results, "11", "9");
      assertExpandGroupCountAndOrder(ddd, 2, results, "12", "14");
    }

    { // queries matching all docs to test null groups from collapse and how it affects expand
      ModifiableSolrParams params = new ModifiableSolrParams();
      params.add("q", "*:*");
      params.add("defType", "edismax");
      params.add("bf", "field(test_i)");
      params.add("expand", "true");
      setDistributedParams(params);

      // nullPolicy=expand
      params.add("fq", "{!collapse field=" + group + " nullPolicy=expand}");

      QueryResponse rsp = query(params);
      assertCountAndOrder(
          6,
          rsp.getResults(),
          "10" /* c */,
          "88" /* null */,
          "2" /* a */,
          "13" /* d */,
          "99" /* null */,
          "6" /* b */);
      Map<String, SolrDocumentList> results = rsp.getExpandedResults();
      assertExpandGroups(results, aaa, bbb, ccc, ddd);
      assertExpandGroupCountAndOrder(aaa, 2, results, "1", "7");
      assertExpandGroupCountAndOrder(bbb, 2, results, "5", "8");
      assertExpandGroupCountAndOrder(ccc, 2, results, "11", "9");
      assertExpandGroupCountAndOrder(ddd, 2, results, "12", "14");

      // nullPolicy=collapse
      params.set("fq", "{!collapse field=" + group + " nullPolicy=collapse}");

      rsp = query(params);
      assertCountAndOrder(
          5,
          rsp.getResults(),
          "10" /* c */,
          "88" /* null */,
          "2" /* a */,
          "13" /* d */,
          "6" /* b */);
      results = rsp.getExpandedResults();
      assertExpandGroups(results, aaa, bbb, ccc, ddd);
      assertExpandGroupCountAndOrder(aaa, 2, results, "1", "7");
      assertExpandGroupCountAndOrder(bbb, 2, results, "5", "8");
      assertExpandGroupCountAndOrder(ccc, 2, results, "11", "9");
      assertExpandGroupCountAndOrder(ddd, 2, results, "12", "14");

      // nullPolicy=collapse w/ expand.nullGroup=true...
      params.set("fq", "{!collapse field=" + group + " nullPolicy=collapse}");
      params.set("expand.nullGroup", "true");

      rsp = query(params);
      assertCountAndOrder(
          5,
          rsp.getResults(),
          "10" /* c */,
          "88" /* null */,
          "2" /* a */,
          "13" /* d */,
          "6" /* b */);
      results = rsp.getExpandedResults();
      assertExpandGroups(results, aaa, bbb, ccc, ddd, null);
      assertExpandGroupCountAndOrder(aaa, 2, results, "1", "7");
      assertExpandGroupCountAndOrder(bbb, 2, results, "5", "8");
      assertExpandGroupCountAndOrder(ccc, 2, results, "11", "9");
      assertExpandGroupCountAndOrder(ddd, 2, results, "12", "14");
      assertExpandGroupCountAndOrder(null, 1, results, "99");

      // nullPolicy=expand w/ expand.nullGroup=true (use small rows to ensure null expanded group)
      params.set("fq", "{!collapse field=" + group + " nullPolicy=expand}");
      params.set("rows", "3");

      rsp = query(params);
      assertCountAndOrder(3, rsp.getResults(), "10" /* c */, "88" /* null */, "2" /* a */);
      results = rsp.getExpandedResults();
      assertExpandGroups(results, aaa, ccc, null);
      assertExpandGroupCountAndOrder(aaa, 2, results, "1", "7");
      assertExpandGroupCountAndOrder(ccc, 2, results, "11", "9");
      assertExpandGroupCountAndOrder(null, 1, results, "99");

      // nullPolicy=expand w/ expand.nullGroup=true & expand.rows = 0
      params.set("expand.rows", "0");

      rsp = query(params);
      assertCountAndOrder(3, rsp.getResults(), "10" /* c */, "88" /* null */, "2" /* a */);
      results = rsp.getExpandedResults();
      assertExpandGroups(results, aaa, ccc, null);
      assertExpandGroupCountAndOrder(aaa, 0, results);
      assertExpandGroupCountAndOrder(ccc, 0, results);
      assertExpandGroupCountAndOrder(null, 0, results);
    }
  }

  private void assertExpandGroups(Map<String, SolrDocumentList> expandedResults, String... groups)
      throws Exception {
    for (String group : groups) {
      if (!expandedResults.containsKey(group)) {
        throw new Exception(
            "Expanded Group Not Found:" + group + ", Found:" + exportGroups(expandedResults));
      }
    }
  }

  private String exportGroups(Map<String, SolrDocumentList> groups) {
    StringBuilder buf = new StringBuilder();
    Iterator<String> it = groups.keySet().iterator();
    while (it.hasNext()) {
      String group = it.next();
      buf.append(group);
      if (it.hasNext()) {
        buf.append(",");
      }
    }
    return buf.toString();
  }

  private void assertExpandGroupCountAndOrder(
      final String group,
      final int count,
      final Map<String, SolrDocumentList> expandedResults,
      final String... docs)
      throws Exception {
    SolrDocumentList results = expandedResults.get(group);
    if (results == null) {
      throw new Exception("Group Not Found:" + group);
    }
    assertCountAndOrder(count, results, docs);
  }

  private void assertCountAndOrder(
      final int count, final SolrDocumentList results, final String... docs) {
    assertEquals(results.toString(), count, results.size());

    for (int i = 0; i < docs.length; i++) {
      String id = docs[i];
      SolrDocument doc = results.get(i);
      assertEquals("Id not in results or out of order", id, doc.getFieldValue("id").toString());
    }
  }
}
